package com.basin.common.core.async;

import com.basin.common.core.async.standard.AbstractBaseTask;
import com.basin.common.core.async.standard.TaskWithResult;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.concurrent.*;

/**
 * @author: Sam ZHONG | sammy1997@aliyun.com
 */
@Slf4j
public class AsyncHelper {

    private static final ExecutorService EXECUTOR = Executors.newVirtualThreadPerTaskExecutor();

    @SafeVarargs
    public static <TASK extends AbstractBaseTask<?>> void executeWithResult(TASK... tasks) {
        List<TASK> taskList = Arrays.asList(tasks);
        invoke(taskList, true);
    }

    public static <TASK extends AbstractBaseTask<?>> void executeWithResult(Collection<TASK> tasks) {
        invoke(tasks, true);
    }

    @SafeVarargs
    public static <TASK extends AbstractBaseTask<?>> void executeWithoutResult(TASK... tasks) {
        List<TASK> taskList = Arrays.asList(tasks);
        invoke(taskList, false);
    }

    public static <TASK extends AbstractBaseTask<?>> void executeWithoutResult(Collection<TASK> tasks) {
        invoke(tasks, false);
    }

    private static <TASK extends AbstractBaseTask<?>> void invoke(Collection<TASK> tasks, boolean needWait) {
        Map<String, TaskWithResult<TASK>> map = new HashMap<>();
        for (TASK task : tasks) {
            Future<?> curFuture = EXECUTOR.submit(task::run);
            map.put(task.getTaskId(), TaskWithResult.<TASK>builder()
                    .future(curFuture)
                    .task(task).build());
        }

        CountDownLatch latch = new CountDownLatch(map.size());

        map.forEach((k, v) -> EXECUTOR.execute(() -> {
            TaskWithResult<TASK> taskTaskWithResult = map.get(k);
            TASK task = taskTaskWithResult.getTask();
            Future<?> future = taskTaskWithResult.getFuture();
            long limitTime = task.getLimitTimeMs();
            try {
                future.get(limitTime, TimeUnit.MILLISECONDS);
                // 成功才执行回调
                EXECUTOR.execute(task::callback);
            } catch (InterruptedException | TimeoutException | ExecutionException e) {
                log.error("获取线程结果异常 {} ", task.getTaskId(), e);
                Thread.currentThread().interrupt();
                future.cancel(true);
            } finally {
                latch.countDown();
            }
        }));

        if (!needWait) return; // 不需要等待结果

        try {
            latch.await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
}
